#pragma once 

#include "ros/ros.h"
#include "ros/serialization.h"
#include "std_msgs/String.h"
#include <vector>
#include <string>
#include <cstring>
#include <chrono>
#include <thread>

class uwb_dt_publisher
{
public:
    typedef struct
    {
        uint8_t topic_index;
        uint8_t id;
    } uwb_dt_param_t;

    enum class frame_type_t
    {
        SINGLE,
        HEAD,
        MID,
        TAIL,
    };

    uwb_dt_publisher() = delete;
    uwb_dt_publisher(ros::NodeHandle &nh)
    {
        data_raw_pub_ = nh.advertise<std_msgs::String>("/nlink_linktrack_data_transmission", 100);
    }
    ~uwb_dt_publisher() = default;

    std::string packData(const frame_type_t& type, const uint32_t& data_size, const uwb_dt_param_t& param, const uint8_t* msg)
    {
        auto head = HEAD_MAP.at(type);
        auto tail = TAIL_MAP.at(type);
        // Frame : head(u16) topic(u8) id(u8) size(u32) data tail(u16)
        uint32_t extra_size = 2 * sizeof(uint16_t) + 2 * sizeof(uint8_t) + sizeof(uint32_t);

        std::vector<uint8_t> buffer(extra_size + data_size, 0);
        buffer[0] = head.first;
        buffer[1] = head.second;
        buffer[2] = param.topic_index;
        buffer[3] = param.id;
        memcpy(&buffer[4], &data_size, sizeof(uint32_t));
        memcpy(&buffer[8], msg, data_size * sizeof(uint8_t));
        buffer[8 + data_size] = tail.first;
        buffer[9 + data_size] = tail.second;
        std::string str;
        str.assign(buffer.begin(), buffer.end());
        return str;
    }

    void publish(const std::string& string_msg)
    {
        std_msgs::String msg;
        msg.data = string_msg;
        data_raw_pub_.publish(msg);
    }

    template <typename T>
    void publish(const uwb_dt_param_t& param, T& msg)
    {
        uint32_t ser_size = ros::serialization::serializationLength(msg);
        std::vector<uint8_t> msg_buffer(ser_size, 0);
        ros::serialization::OStream stream(msg_buffer.data(), ser_size);
        ros::serialization::serialize(stream, msg);

        uint32_t msg_buffer_size = ser_size;
        uint32_t curr_pointer = 0;
        do {
            uint32_t pub_size = (msg_buffer_size > FRAME_MAX_SIZE) ? FRAME_MAX_SIZE : msg_buffer_size;
            frame_type_t frame_type = (curr_pointer == 0) 
                                    ? (msg_buffer_size > FRAME_MAX_SIZE) ? frame_type_t::HEAD : frame_type_t::SINGLE
                                    : (msg_buffer_size > FRAME_MAX_SIZE) ? frame_type_t::MID : frame_type_t::TAIL;
            publish(packData(frame_type, pub_size, param, &msg_buffer[curr_pointer]));

            msg_buffer_size -= pub_size;
            curr_pointer += pub_size;

            std::this_thread::sleep_for(std::chrono::milliseconds(20));
        } while (curr_pointer < ser_size);
    }

    template <typename T>
    static void unpackData(int& id, T& receive, const std_msgs::String::ConstPtr& msg)
    {
        const std::string &data = msg->data;
        std::vector<uint8_t> buffer;
        buffer.assign(msg->data.begin(), msg->data.end());  // deep copy

        id = int(buffer[1]);
        uint32_t ser_size;
        memcpy(&ser_size, &buffer[2], sizeof(uint32_t));
        uint8_t *ptr = &buffer[6];
        ros::serialization::IStream stream(ptr, ser_size);
        ros::serialization::Serializer<T>::read(stream, receive);
    }
private:
    ros::Publisher data_raw_pub_;

public:
    // 200 Bytes is the maximum size of a frame
    constexpr static uint32_t FRAME_MAX_SIZE = 180;

    constexpr static uint8_t FRAME_HEAD = 0x55;
    constexpr static uint8_t FRAME_PAUSE = 0xA1;
    constexpr static uint8_t FRAME_RESUME = 0xA2;
    constexpr static uint8_t FRAME_TAIL = 0x04;
    const std::map<frame_type_t, std::pair<uint8_t, uint8_t>> HEAD_MAP = {
        {frame_type_t::SINGLE,  {FRAME_HEAD, FRAME_HEAD}},
        {frame_type_t::HEAD,    {FRAME_HEAD, FRAME_HEAD}},
        {frame_type_t::MID,     {FRAME_HEAD, FRAME_RESUME}},
        {frame_type_t::TAIL,    {FRAME_HEAD, FRAME_RESUME}},
    };
    const std::map<frame_type_t, std::pair<uint8_t, uint8_t>> TAIL_MAP = {
        {frame_type_t::SINGLE, {FRAME_TAIL, FRAME_TAIL}},
        {frame_type_t::HEAD,   {FRAME_PAUSE, FRAME_TAIL}},
        {frame_type_t::MID,    {FRAME_PAUSE, FRAME_TAIL}},
        {frame_type_t::TAIL,   {FRAME_TAIL, FRAME_TAIL}},
    };
};